package udpecho

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"github.com/kataras/golog"
	"nprobe/core"
	"nprobe/enviroments"
	"nprobe/utils/routines"
	"time"
)

type UDPEchoReceiver struct {
	core.BaseReceiver
}

func NewUDPEchoReceiver(
	ctx context.Context, name string, id int, task *core.ProbeTask, salts map[int][]byte, targetSet map[string]bool,
) *UDPEchoReceiver {
	ppRBuf := core.NewPPRBuf(task.SpanMs, task.TimeoutMs)

	return &UDPEchoReceiver{
		BaseReceiver: core.BaseReceiver{
			Ctx:      ctx,
			Name:     name,
			ID:       id,
			Task:     task,
			RBuf:     ppRBuf,
			Salts:    salts,
			DstIPSet: targetSet,
		},
	}
}

func (r *UDPEchoReceiver) StartReceiver() error {
	routines.GoWithRecover(r.Name, r.doReceiveUDP)
	return nil
}

func (r *UDPEchoReceiver) doReceiveUDP() {
	conn, err := core.OpenReceiverIP4PacketConn("ip4:udp", "0.0.0.0") // 本地地址
	if err != nil {
		return
	}
	r.Conn = conn

	pid := enviroments.PID
	filter := core.CreateBPFUDP(uint16(pid), uint16(r.ID), r.Task.SrcPortMin, r.Task.SrcPortMax, r.Task.DstPortMin, r.Task.DstPortMax)
	err = core.SetRawConnBPF(conn, filter)
	if err != nil {
		fmt.Println("unable to open network connection")
		return
	}

	r.CurrSpanRecord = core.NewSpanRecord(0)
	payloadBuf := core.AllocRxBuf(r.Task.MsgLen)
	for {
		select {
		case <-r.Ctx.Done():
			r.StopReceiver()
			return
		default:
			r.receiveOneUDPPkt(payloadBuf)
		}

	}
}

func (r *UDPEchoReceiver) receiveOneUDPPkt(payloadBuf []byte) {

	conn := r.Conn
	err := conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	if err != nil {
		return
	}

	h, payload, _, err := conn.ReadFrom(payloadBuf)
	if err != nil {
		return
	}

	rxTs := time.Now().UnixNano()
	if h.Version != 4 {
		return
	}

	remoteIP := h.Src.To4().String()
	localIP := h.Dst.To4().String()
	if !r.DstIPSet[remoteIP] {
		return
	}

	r.processUDPEchoReply(localIP, remoteIP, rxTs, payload)
}

func (r *UDPEchoReceiver) processUDPEchoReply(localIP string, remoteIP string, rxTs int64, payload []byte) {
	// udp header
	srcPort := binary.BigEndian.Uint16(payload[0:2])
	dstPort := binary.BigEndian.Uint16(payload[2:4])

	// udp payload
	udpLoad := payload[8:]
	txTs := int64(binary.BigEndian.Uint64(udpLoad[12:20]))
	seq := binary.BigEndian.Uint16(udpLoad[20:22])

	bitFlip := false
	if r.Task.BitFlip {
		bitFlip = !bytes.Equal(r.Salts[int(seq)%4][22:], udpLoad[22:])
		if bitFlip {
			golog.Debugf("salt: %x", r.Salts[int(seq)%4])
			golog.Debugf("data: %x", payload)
			golog.Debugf("[bitflip] src=%s  dst=%s seq=%d, ts=%d", localIP, remoteIP, seq, txTs)
		}
	}

	r.recordPkt(localIP, remoteIP, dstPort, srcPort, seq, txTs, rxTs, bitFlip)
}

func (r *UDPEchoReceiver) recordPkt(localIP string, remoteIP string, srcPort uint16, dstPort uint16, seq uint16, txTs int64, rxTs int64, bitFlip bool) {
	spanID := txTs / (r.Task.SpanMs * int64(time.Millisecond))
	golog.Debugf("%s: %s:%d->%s:%d, seq: %d, spanID: %d, rxTs: %d, txTs: %d, delta: %d, bitFlip: %t",
		r.Name, localIP, srcPort, remoteIP, dstPort, seq, spanID, rxTs, txTs, rxTs-txTs, bitFlip)
	if spanID > r.CurrSpanRecord.SpanID {
		r.CurrSpanRecord = core.NewSpanRecord(spanID)
		r.RBuf.Enqueue(r.CurrSpanRecord)
	}

	var updateSpanRecord *core.SpanRecord
	var offset uint64
	for {
		spanRecord, err := r.RBuf.PeekPing(offset)
		if err != nil {
			break
		}

		if spanID == spanRecord.SpanID {
			updateSpanRecord = spanRecord
			break
		}

		offset++
	}

	if updateSpanRecord != nil {
		key := fmt.Sprintf("%s-%d", remoteIP, seq)
		updateSpanRecord.PktRecordMap[key] = &core.PktRecord{
			Key:      key,
			SrcIP:    localIP,
			DstIP:    remoteIP,
			SrcPort:  srcPort,
			DstPort:  dstPort,
			Proto:    r.Task.Proto,
			Tos:      r.Task.TOS,
			RxTxTs:   txTs,
			RxTs:     rxTs,
			Received: true,
			BitFlip:  bitFlip,
		}
	}
}
